#include "refine/smooth_c.h"
#include "refine/window_pass.h"
#include "refine/varmap.h"
#include <limits>
#include <iostream>
#include <set>
#include "imglib/imageclass.h"
#include <imglib/bitmap.h>
#include <imglib/bmpmorph.h>

#define GSL_RANGE_CHECK_OFF
#define HAVE_INLINE
#include <gsl/gsl_linalg.h>

using std::clog;
using std::endl;
using std::flush;

namespace refine {

namespace smooth_c {

const params default_params = {2, 2, 1, 100, 100, false};

typedef grey_mask::mask::elem_type grey_mask_value;
const grey_mask_value grey_mask_max_value = std::numeric_limits<grey_mask_value>::max();	//255
const grey_mask_value grey_mask_min_value = std::numeric_limits<grey_mask_value>::min();	//0
const grey_mask_value grey_mask_avg_value = (grey_mask_max_value + grey_mask_min_value) / 2;//127

typedef unsigned uint;

template <typename T>
inline T sqr(T t)
	{ return t * t; }

/// linear relation of central pixel's alpha is found in the form:
/// alpha = 
///   r[0] * sum_i R_i a_i +
///   r[1] * sum_i G_i a_i +
///   r[2] * sum_i B_i a_i +
///   r[3] * sum_i X_i a_i +
///   r[4] * sum_i Y_i a_i +
///   r[5] * sum_i a_i
struct relation {
	float v[6];
	float operator[] (uint i) const { return v[i]; }
	float & operator[] (uint i) { return v[i]; }
};
typedef std::vector<relation> relations;

class relation_constructor
{
public:
	relation_constructor(byte_image const & im_, 
		varmap const & vmap_, relations & rels_)
		: im(im_), vmap(vmap_), rels(rels_)
	{ 
		rels.resize(vmap.count());
		remove_all(); 
		constructions_count = 0; 
		Av = gsl_matrix_view_array(&A[0][0], 6, 6);
		Vv = gsl_matrix_view_array(&V[0][0], 6, 6);
		Sv = gsl_vector_view_array(S, 6);
		work_v = gsl_vector_view_array(work, 6);
	}
	size_extent dim() const
		{ return im.dim(); }
	uint count() const
		{ return constructions_count; }
	bool selected(size_point const & p) const
		{ return vmap.variable(p); }
	void add(size_point const & p) {
		byte_color const & I = im[p];
		add_sample(I.r, I.g, I.b, p.x, p.y);
	}
	void remove(size_point const & p) {
		byte_color const & I = im[p];
		remove_sample(I.r, I.g, I.b, p.x, p.y);
	}
	void remove_all() {
		n = 0;
		R = G = B = X = Y = 0;
		RR = RG = RB = RX = RY = GG = GB = GX = GY = BB = BX = BY = XX = XY = YY = 0;
	}
	void process(size_point const & p, size_rect const &) {
		++constructions_count;
		byte_color const & I = im[p];
		process_core(rels[vmap[p]], I.r, I.g, I.b, p.x, p.y);
	}
private:
	void add_sample(uint r, uint g, uint b, uint x, uint y) {
		++n;
		R += r; 
		G += g; 
		B += b; 
		X += x; 
		Y += y; 
		RR += r * r;
		RG += r * g;
		RB += r * b;
		RX += r * x;
		RY += r * y;
		GG += g * g;
		GB += g * b;
		GX += g * x;
		GY += g * y;
		BB += b * b;
		BX += b * x;
		BY += b * y;
		XX += x * x;
		XY += x * y;
		YY += y * y;
	}
	void remove_sample(uint r, uint g, uint b, uint x, uint y) {
		--n;
		R -= r; 
		G -= g; 
		B -= b; 
		X -= x; 
		Y -= y; 
		RR -= r * r;
		RG -= r * g;
		RB -= r * b;
		RX -= r * x;
		RY -= r * y;
		GG -= g * g;
		GB -= g * b;
		GX -= g * x;
		GY -= g * y;
		BB -= b * b;
		BX -= b * x;
		BY -= b * y;
		XX -= x * x;
		XY -= x * y;
		YY -= y * y;
	}
	void process_core(relation & rel, uint rc, uint gc, uint bc, uint xc, uint yc) {
		//fill A matrix with values like (A-Ac)(B-Bc)
		A[0][0] = RR + rc*rc*n - 2*R*rc;
		A[0][1] = A[1][0] = int(RG + rc*gc*n - R*gc - G*rc);
		A[0][2] = A[2][0] = int(RB + rc*bc*n - R*bc - B*rc);
		A[0][3] = A[3][0] = int(RX + rc*xc*n - R*xc - X*rc);
		A[0][4] = A[4][0] = int(RY + rc*yc*n - R*yc - Y*rc);
		A[0][5] = A[5][0] = int(R - rc*n);
		A[1][1] = GG + gc*gc*n - 2*G*gc;
		A[1][2] = A[2][1] = int(GB + gc*bc*n - G*bc - B*gc);
		A[1][3] = A[3][1] = int(GX + gc*xc*n - G*xc - X*gc);
		A[1][4] = A[4][1] = int(GY + gc*yc*n - G*yc - Y*gc);
		A[1][5] = A[5][1] = int(G - gc*n);
		A[2][2] = BB + bc*bc*n - 2*B*bc;
		A[2][3] = A[3][2] = int(BX + bc*xc*n - B*xc - X*bc);
		A[2][4] = A[4][2] = int(BY + bc*yc*n - B*yc - Y*bc);
		A[2][5] = A[5][2] = int(B - bc*n);
		A[3][3] = XX + xc*xc*n - 2*X*xc;
		A[3][4] = A[4][3] = int(XY + xc*yc*n - X*yc - Y*xc);
		A[3][5] = A[5][3] = int(X - xc*n);
		A[4][4] = YY + yc*yc*n - 2*Y*yc;
		A[4][5] = A[5][4] = int(Y - yc*n);
		A[5][5] = n;

		//compute: A = U * S * V^T; U replaces A
		gsl_linalg_SV_decomp(&Av.matrix, &Vv.matrix, &Sv.vector, &work_v.vector);

		//work := S^-1 * last_row_of_V
		for (uint i = 0; i < 6; ++i) {
			if (S[i] < std::numeric_limits<double>::epsilon())
				work[i] = 0;
			else
				work[i] = V[5][i] / S[i];
		}

		//relation := U * work
		for (uint i = 0; i < 6; ++i) {
			double s = 0;
			for (uint j = 0; j < 6; ++j)
				s += A[i][j] * work[j];
			rel[i] = s;
		}

		//when we will use relation, sum(R_i a_i) will be substituted instead of sum((R_i - rc) a_i)
		rel[5] -= rel[0] * rc + rel[1] * gc + rel[2] * bc + rel[3] * xc + rel[4] * yc;
		//float one = R * rel[0] + G * rel[1] + B * rel[2] + X * rel[3] + Y * rel[4] + n * rel[5];
		//if (one < 0.75f || one > 1.25f)
		//	assert (false);
	}
private:
	byte_image const & im;
	varmap const & vmap;
	relations & rels;
	uint constructions_count;
private:
	uint n;
	uint R, G, B, X, Y;
	uint RR, RG, RB, RX, RY, GG, GB, GX, GY, BB, BX, BY, XX, XY, YY;
private:
	double A[6][6], V[6][6], S[6], work[6];
	gsl_matrix_view Av, Vv;
	gsl_vector_view Sv, work_v;
};

//template <typename Raster, typename Matte>
class smoother
{
public:
	smoother(byte_image const & src_, varmap const & vmap_, relations const & rels_,
		grey_mask::mask const & matte_, grey_mask::mask & next_matte_,
		img::bitmap & modifiable_, float max_condition_number_)
		: src(src_), vmap(vmap_), rels(rels_), matte(matte_), next_matte(next_matte_), modifiable(modifiable_),
		  max_condition_number(max_condition_number_)
		{ clear(); }
	void clear() {
		checks = achanges = 0;
		remove_all();
	}
	size_extent dim() const
		{ return src.dim(); }
	uint check_count() const
		{ return checks; }
	uint change_count() const
		{ return achanges; }
	bool selected(size_point const & p) const
		{ return modifiable[p]; }
	void add(size_point const & p) {
		byte_color const & I = src[p];
		add_sample(I.r, I.g, I.b, p.x, p.y, matte[p]);
	}
	void remove(size_point const & p) {
		byte_color const & I = src[p];
		remove_sample(I.r, I.g, I.b, p.x, p.y, matte[p]);
	}
	void remove_all() 
		{ n = RA = GA = BA = XA = YA = A = AA = 0; }
	///checks max_condition_number and returns true in the case of fail
	bool bad_condition_number() const {
		if (n == 0)
			return true;
		double a = AA;
		double b = grey_mask_max_value * A - AA;
		double c = grey_mask_max_value * (grey_mask_max_value * n - 2 * A) + AA;
		double D = sqr(a - c) + 4 * b * b;
		return sqr(max_condition_number + 1) * D > sqr((max_condition_number - 1) * (a + c));
	}
	void process(size_point const & p, size_rect const &) {
		++checks;
		if (bad_condition_number()) {
			modifiable[p] = false;
			return;
		}

		relation const & r = rels[vmap[p]];
		float a = r[0] * RA + r[1] * GA + r[2] * BA + r[3] * XA + r[4] * YA + r[5] * A + 0.5f; //0.5 for proper rounding

		grey_mask_value g;
		if (a <= grey_mask_min_value)
			g = grey_mask_min_value;
		else if (a >= grey_mask_max_value)
			g = grey_mask_max_value;
		else 
			g = grey_mask_value(a);

		if (matte[p] == g)
			modifiable[p] = false;
		else {
			++achanges;
			next_matte[p] = g;
		}
	}
private:
	void add_sample(uint r, uint g, uint b, uint x, uint y, uint a) {
		++n;
		RA += r * a;
		GA += g * a;
		BA += b * a;
		XA += x * a;
		YA += y * a;
		A += a;
		AA += a * a;
	}
	void remove_sample(uint r, uint g, uint b, uint x, uint y, uint a) {
		--n;
		RA -= r * a;
		GA -= g * a;
		BA -= b * a;
		XA -= x * a;
		YA -= y * a;
		A -= a;
		AA -= a * a;
	}
private:
	byte_image const & src;
	varmap const & vmap;
	relations const & rels;
	grey_mask::mask const & matte;
	grey_mask::mask & next_matte;
	img::bitmap & modifiable;
	float max_condition_number;
	uint checks, achanges;
private:
	uint n, RA, GA, BA, XA, YA, A, AA;
};

void perform_core(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const img::bitmap & strict_fore, 
	const img::bitmap & strict_back,
	img::bitmap & modifiable, ///< map of pixels that may be changed during the following pass
	const params & p)
{
	if (!p.silent)
		clog << "Spatial Volosok(R) matting..." << endl;

	size_extent dim = src.dim();
	assert (dim == matte.dim());
	assert (dim == modifiable.dim());
	int_rect frame(0, 0, dim.sx, dim.sy);

	varmap vmap(strict_fore, strict_back);
	relations rels(vmap.count());
	window_pass(relation_constructor(src, vmap, rels), p.whs);

	grey_mask::mask next_matte = matte;

	for (uint pass = 0; pass < p.passes_num; ++pass) {
		if (!p.silent)
			clog << "pass " << pass+1 << '/' << p.passes_num << flush;

		img::combine(modifiable, strict_fore, img::bit_and_not());
		img::combine(modifiable, strict_back, img::bit_and_not());

		smoother sm(src, vmap, rels, matte, next_matte, modifiable, p.max_condition_number);
		window_pass(sm, p.whs);
		if (!p.silent)
			clog << ": checked = " << sm.check_count() << ", alpha changes = " << sm.change_count() << endl;
		if (sm.change_count() == 0)
			break;

		matte = next_matte;

		img::dilate8(modifiable, p.whs);
	}
}


void perform(
	const byte_image & src, 
	const bit_mask::mask & m, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = src.dim();
	assert (dim == m.dim());

	matte.resize(dim);
	for (grey_mask::mask::iterator i = matte.begin(); i != matte.end(); ++i)
		*i = m[i] ? grey_mask_max_value : grey_mask_min_value;

	img::bitmap strict_fore, strict_back;
	img::erode4(strict_fore, m, p.fore_depth);
	img::dilate4(strict_back, m, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable = m; 
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, w);
		img::bitmap tmp = m;
		tmp.inverse();
		img::dilate4(tmp, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	const byte_image & src, 
	grey_mask::mask & matte, 
	const params & p)
{
	size_extent dim = matte.dim();

	img::bitmap not_min(dim), not_max(dim), strict_fore(dim), strict_back(dim);
	for (grey_mask::mask::const_iterator i = matte.begin(); i != matte.end(); ++i) {
		not_min[i] = *i != grey_mask_min_value;
		not_max[i] = *i != grey_mask_max_value;
		strict_fore[i] = *i > grey_mask_avg_value;
	}

	strict_back = strict_fore;
	img::erode4(strict_fore, p.fore_depth);
	img::dilate4(strict_back, p.back_depth);
	strict_back.inverse();

	img::bitmap modifiable;
	{
		uint w = 1 + p.whs * (p.max_condition_number - 1) / (p.max_condition_number + 1);
		img::dilate4(modifiable, not_min, w);
		img::bitmap tmp;
		img::dilate4(tmp, not_max, w);
		img::combine(modifiable, tmp, img::bit_and());
	}

	perform_core(src, matte, strict_fore, strict_back, modifiable, p);
}

void perform(
	argb_byte_image & src,
	const params & p)
{
	//!!! temporary solution
	byte_image src2 = src; 
	grey_mask::mask matte = img::projection(src, &argb_byte_color::a);

	perform(src2, matte, p);

	img::projection(src, &argb_byte_color::a) = matte;
}

} //namespace smooth_c

} //namespace refine
